#import tensorflow as tf
from ais_bench.infer.interface import InferSession
import numpy as np
import soundfile as sf

if __name__ == "__main__":
    om = InferSession(0, '../fastspeech2/OM/fs2_model_linux_x86_64.om', debug=False)
    index = 3
    if index == 0:
        input_phone_ids = np.array([108,67,75,68,88,6,34,20,30,18,35,3,34,6,34,1,30,15,39,22,57,3,26,80,72,87,72,98,107], dtype= np.int32)
        input_tone_ids = np.array([0,7,0,6,0,0,2,0,1,0,3,0,5,0,2,0,4,0,4,0,4,0,4,0,7,0,8,0,0], dtype= np.int32)
        input_prosody_ids = np.array([4,0,0,0,3,0,2,0,1,0,1,0,2,0,1,0,3,0,2,0,1,0,2,0,0,0,0,4,4], dtype= np.int32)
    elif index == 1:
        input_phone_ids = np.array([108, 9, 47, 3, 26, 18, 51, 0, 51, 12, 48, 4, 29, 3, 34, 102, 103, 86, 7, 44, 20, 34, 21, 30, 8, 27, 15, 39, 9, 34, 107], dtype= np.int32)
        input_tone_ids = np.array([0, 0, 2, 0, 4, 0, 4, 0, 4, 0, 2, 0, 2, 0, 5, 0, 7, 0, 0, 4, 0, 4, 0, 4, 0, 1, 0, 3, 0, 5, 0], dtype= np.int32)
        input_prosody_ids = np.array([4, 0, 2, 0, 1, 0, 2, 0, 1, 0, 1, 0, 1, 0, 2, 0, 0, 3, 0, 1, 0, 1, 0, 2, 0, 1, 0, 1, 0, 4, 4], dtype= np.int32)
    elif index == 2:
        input_phone_ids = np.array([108,  10,  35,   5,  34,  12,  27,  19,  39,   3,  32,  21,  32,  22,  39,   7,  39,  20, 54,  10,  37,   3,  34,  15,  34,   6,  57,  66,  73,  87,  68,  79,  77,  97,  68,  89, 106,  17,  26,  10,  37,  17, 33,   5,  56,   6,  34,  22,  56,   2,  51,   7,  47,  20, 34,   5,  34,  15,  39,   7,  46,   3,  34,   7,  47,  0,  51, 107], dtype= np.int32)
        input_tone_ids = np.array([0, 0, 3, 0, 4, 0, 4, 0, 4, 0, 1, 0, 3, 0, 4, 0, 3, 0, 1, 0, 2, 0, 5, 0, 4, 0, 4, 8, 0, 0, 6, 0, 7, 0, 6, 0, 0, 0, 1, 0, 5, 0, 1, 0, 4, 0, 2, 0, 4, 0, 4, 0, 4, 0, 4, 0, 4, 0, 4, 0, 4, 0, 5, 0, 4, 0, 4, 0], dtype= np.int32)
        input_prosody_ids = np.array( [4, 0, 1, 0, 2, 0, 1, 0, 2, 0, 1, 0, 2, 0, 1, 0, 2, 0, 1, 0, 1, 0, 2, 0, 1, 0, 2, 0, 0, 0, 0, 0, 0, 0, 0, 4, 4, 0, 1, 0, 2, 0, 1, 0, 2, 0, 1, 0, 3, 0, 1, 0, 2, 0, 1, 0, 2, 0, 1, 0, 1, 0, 2, 0, 1, 0, 4, 4], dtype= np.int32)
    elif index == 3:
        input_phone_ids = np.array([108,  20,  33,  21,  30,   7,  39,  18,  35,   5,  56,   7,  40,   7,  42,  18,  35, 18, 30,  20,  29,   4,  26,   0,  51,  19, 41,  19, 39, 107], dtype= np.int32)
        input_tone_ids = np.array([0, 0, 1, 0, 1, 0, 4, 0, 3, 0, 2, 0, 1, 0, 1, 0, 3, 0, 3, 0, 4, 0, 1, 0, 4, 0, 1, 0, 1, 0], dtype= np.int32)
        input_prosody_ids = np.array([4, 0, 1, 0, 1, 0, 1, 0, 2, 0, 1, 0, 2, 0, 1, 0, 2, 0, 1, 0, 2, 0, 1, 0, 2, 0, 1, 0, 4, 4], dtype= np.int32)
    input_phone_ids   = np.expand_dims(input_phone_ids, 0)
    input_tone_ids    = np.expand_dims(input_tone_ids, 0)
    input_prosody_ids = np.expand_dims(input_prosody_ids, 0)

    speed_ratios=np.array([1.0], dtype=np.float32)
    f0_ratios =np.array([1.0], dtype=np.float32)
    energy_ratios =np.array([1.0], dtype=np.float32)

    output = om.infer([input_prosody_ids, input_phone_ids,speed_ratios, input_tone_ids,f0_ratios,energy_ratios], "dymshape", [2147583648,2147583648,1000,2147583648,1000,1000])
    [mel_before, mel_after, duration_outputs, f0_outputs, energy_outputs] = output
    print(mel_before)
    mel_shape = mel_before.shape
    bs = mel_shape[0]
    mel_len = mel_shape[1]
    print("mel_before shape {}".format(mel_shape))
    print('fs2.output:---------------')

    generator_om = InferSession(1, '../acl_hifi_gan/OM/dynamic_dims_new_linux_x86_64.om')
    custom_sizes = bs * mel_len * 300 * 4
    wavs = generator_om.infer([mel_before], mode='dymshape', custom_sizes=custom_sizes)
    print(wavs[0].shape)
    wavs = wavs[0][:, :, 0]
    print("wavs shape {}".format(wavs.shape))
    sf.write('./test_{}.wav'.format(index), wavs.T, 24000, "PCM_16")
